using System.Text;
using System.Text.RegularExpressions;
using Crux.Core.LangChain.Abstractions.DocumentLoaders;
using Tabula;
using Tabula.Detectors;
using Tabula.Extractors;
using UglyToad.PdfPig;
using UglyToad.PdfPig.Content;
using UglyToad.PdfPig.Core;

namespace Crux.Core.LangChain.DocumentLoaders.Pdf;

/// <summary>
/// 
/// </summary>
public class PdfLoader : IDocumentLoader
{
    private readonly static IExtractionAlgorithm _extractionAlgorithm = new SpreadsheetExtractionAlgorithm();
    private readonly static IDetectionAlgorithm _detectionAlgorithm = new SpreadsheetDetectionAlgorithm();

    /// <summary>
    /// 
    /// </summary>
    /// <param name="dataSource"></param>
    /// <param name="settings"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    /// <exception cref="ArgumentNullException"></exception>
    public async Task<IReadOnlyCollection<Document>> LoadAsync(
        DataSource dataSource,
        DocumentLoaderSettings? settings = null,
        CancellationToken cancellationToken = default)
    {
        dataSource = dataSource ?? throw new ArgumentNullException(paramName: nameof(dataSource));

        using var stream = await dataSource.GetStreamAsync(cancellationToken).ConfigureAwait(false);

        using var document = PdfDocument.Open(stream, new ParsingOptions());

        var metadata = settings.CollectMetadataIfRequired(dataSource);

        return document
            .GetPages()
            .Select(page => new Document(
                ExtractText(page),
                metadata?.With(new Dictionary<string, object>
                {
                    { "page", page.Number },
                })))
            .Where(z => !string.IsNullOrEmpty(z.PageContent))
            .ToArray();
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="page"></param>
    /// <returns></returns>
    private static string ExtractText(Page page)
    {
        var tableAreas = page.GetTablesAreas(_detectionAlgorithm);
        var words = page.GetWords();

        // 预估StringBuilder的初始容量，减少重新分配
        var estimatedCapacity = words.Sum(w => w.Text.Length) + tableAreas.Count() * 100;
        var sb = new StringBuilder(estimatedCapacity);

        var contentElements = new List<(PdfRectangle Area, object Element)>();

        // 处理文字元素
        foreach (var word in words)
        {
            if (!tableAreas.Any(x => x.Top > word.BoundingBox.Top && x.Bottom < word.BoundingBox.Bottom))
            {
                contentElements.Add((word.BoundingBox, word));
            }
        }

        // 添加表格区域
        contentElements.AddRange(tableAreas.Select(area => (area, (object)area)));

        // 按照从上到下的顺序排序
        contentElements.Sort((a, b) => b.Area.Bottom.CompareTo(a.Area.Bottom));

        var previous = default(Letter);

        foreach (var (area, element) in contentElements)
        {
            if (element is Word word)
            {
                ProcessWordElement(page, word, area, ref previous, sb);
            }
            else if (element is PdfRectangle rectangle)
            {
                ProcessTableElement(page, area, sb);
            }
        }

        contentElements.Clear();
        contentElements.Capacity = 1;

        contentElements = null;

        return Regex.Replace(sb.ToString(), @"^\s+$[\r\n]*", "", RegexOptions.Multiline);
    }

    private static void ProcessWordElement(
        in Page page,
        in Word word,
        in PdfRectangle area,
        ref Letter? previous,
        StringBuilder sb)
    {
        var letter = word.Letters?.FirstOrDefault();
        if (letter is not null)
        {
            if (IsNewline(previous, letter, page, out var isDoubleNewline))
            {
                if (previous?.Value == " ")
                {
                    sb.Remove(sb.Length - 1, 1);
                }

                sb.AppendLine();
            }
        }

        previous = word.Letters?.LastOrDefault();
        sb.Append(word.Text);
    }

    static bool IsNewline(Letter? previous, Letter letter, Page page, out bool isDoubleNewline)
    {
        isDoubleNewline = false;

        if (previous == null)
        {
            return false;
        }

        var ptSizePrevious = (int)Math.Round(previous.PointSize);
        var ptSize = (int)Math.Round(letter.PointSize);
        var minPtSize = ptSize < ptSizePrevious ? ptSize : ptSizePrevious;

        var gap = Math.Abs(previous.StartBaseLine.Y - letter.StartBaseLine.Y);

        if (gap > minPtSize * 1.7 && previous.StartBaseLine.Y > letter.StartBaseLine.Y)
        {
            isDoubleNewline = true;
        }

        return gap > minPtSize * 0.9;
    }

    private static void ProcessTableElement(
        in Page page,
        in PdfRectangle area,
        in StringBuilder sb)
    {
        var range = new PdfRectangle(
            area.TopLeft,
            area.TopRight,
            area.BottomLeft.MoveY(-10),
            area.BottomRight.MoveY(-10));

        var tables = page.GetTables([range], _extractionAlgorithm);

        sb.AppendLine();
        foreach (var table in tables)
        {
            FormatTableContent(table, sb);
        }
    }

    private static void FormatTableContent(
        in Table table, in StringBuilder sb)
    {
        var builder = new StringBuilder();
        var isFirstRow = true;

        foreach (var row in table.Rows)
        {
            if (isFirstRow)
            {
                isFirstRow = false;
                builder.AppendLine("| " + string.Join(" | ", row.Select(x => x.GetText().Replace("\r", ""))) + " |");
                builder.AppendLine("| " + string.Join(" | ", row.Select(_ => "-----------")) + " |");
            }
            else
            {
                builder.AppendLine("| " + string.Join(" | ", row.Select(x => x.GetText().Replace("\r", ""))) + " |");
            }
        }

        var tableString = builder.ToString();

        sb.AppendLine(builder.ToString());
    }
}